#from openopt import NLP, GLP
from scipy import optimize
from numpy import eye, cos, arange, ones, asarray, abs, zeros, sqrt, asscalar, array, radians, degrees
import pylab
from string import rjust, ljust, expandtabs
import time
import JT


class Test():
    def __init__(self):            
        #===============================================================================
        #    Data Paths
        #===============================================================================
        fixedImageName = 'C:/Users/bryan/bryan-code/2D3D/vert1/fluoro/ushortim080-LAT.mhd'
        #fixedImageName = 'C:/Users/bryan/bryan-code/trunk/Images/CalibratedDRRImage.mhd'
        inputVolFileName = 'C:/Users/bryan/bryan-code/2D3D/vert1/CT/ucharCTv1r1.mha'
        staFile = 'C:/Users/bryan/bryan-code/2D3D/vert1/CT/ucharCTv1r1.sta'
        calFileLat = 'C:/Users/bryan/bryan-code/trunk/TestData/ext_cal1.txt'
        
        roi = [[210,187],[277,370]]
        
        #===============================================================================
        #    Setup DRR
        #===============================================================================
        xraycam = JT.XRayCamera()
        drr = JT.DRR()
        drr.SetXRayCamera(xraycam)
        
        drr.SetBlackOnWhiteImage()
        drr.InteractionOff()
        
        cal = JT.ExternalCalibration()
        cal.LoadConsolidatedCalibration(calFileLat)
        cal.LoadStaFile(staFile)
        
        xraycam.SetExternalCalibration(cal)
        
        volume = JT.Volume()
        volume.SetVolumeFilename(inputVolFileName)
        volume.SetOriginalTransform(cal._VolumeT)
        volume.SetBodyCenter((3.906, 2.17, -7.378))
        drr.AddVolume(volume)
        
        # Set Fixed image using path
        fixedImage = JT.FixedImage()
        fixedImage.SetFileName(fixedImageName)
        
        self.reg=JT.Registration()
        self.reg.SetFixedImage(fixedImage.GetImage(0))
        self.reg.SetMovingImageGenerator(drr)
        self.reg.SetRegionOfInterest(roi)
        
        self.T = JT.Testing()


    def run_test(self,n_pose=50,method='bpc',debug=True,plot=False):    
    #===============================================================================
    #    Setup Optimizer
    #===============================================================================
    # Select the desired starting pose, between [0,200)
        pose = -self.T.GetStartPose(n_pose)
        print "Starting Pose: ", pose
        
        # Choose metric and bounds (step size?) for search
        self.reg.SetImageMetric(method)

        # Optimizer specific parameters

        #x0 = pose
        def my_func(x):
            x[3:6] = radians(x[3:6])
            #x_opt[3:6]=x
            val=self.reg.GetValue(x,debug=True,plot=False)
            return val
        f = my_func
        #f = self.reg.GetValue
        #For NCC metric, T0 should be 0.35999
        T0=1.0
        Tf=1e-20
        maxiter=10
        maxeval=200
        maxaccept=100
        dwell=20
        feps=1e-10
        
        print "Phase I..."
        start_time = time.time()
        #x0_1=pose[3:6]
        x0=asarray(pose)
        x0[3:6]=degrees(pose[3:6])
        step_length1 = eye(6)/100
        retval1 = optimize.fmin_powell(
                                     f,x0=x0,
                                     xtol = 0.01, ftol = 0.1, 
                                     maxiter = 10, maxfun = 10, 
                                     full_output = 1, disp = 1, retall = 1, 
                                     callback = None, direc = step_length1)
        (xopt1,fopt1,direc1,iter1,funcalls1,warnflag1,allvecs1) = retval1
        print xopt1
        
#        print "Finished Phase I"
#        print "Phase II..."
#        f = self.reg.GetValue
#        x0_2 = pose.copy()
#        x0_2[3:6] = xopt1
#        step_length2 = eye(6)/10
#        step_length2[3][3] = step_length2[4][4] = step_length2[5][5] = step_length2[0][0]/10
#        
#        retval2 = optimize.fmin_powell(
#                                     f, x0_2, args=(debug,plot),
#                                     xtol = 0.0001, ftol = 0.001, 
#                                     maxiter = 10, maxfun = 100, 
#                                     full_output = 1, disp = 1, retall = 1, 
#                                     callback = None, direc = step_length2)
#        (xopt,fopt,direc,iter,funcalls,warnflag,allvecs) = retval2
        end_time = time.time()
        
        #===============================================================================
        #    Display/Save Results
        #===============================================================================
        #(params, fopt, direc, numiter, func_calls, warnflag, allvecs) = retval
        print "Finished Phase II"
#        print "Performed %i function evaluations, in %f seconds, best cost function: %f" \
#               % (feval,end_time-start_time,Jmin)
#        
        print "Starting Pose: ", pose
        print "Final Pose   : ", xopt1
        
        # Save registration data to file for analysis of tuning parameters
#        log_path = 'C:/Users/bryan/bryan-code/trunk/Results/Reg.log'
#        log = open(log_path,'a')
#        output = (str(Jmin),str(xmin[0]),str(xmin[1]),str(xmin[2]),str(xmin[3]),str(xmin[4]),str(xmin[5]),\
#                  str(x0[0]),str(x0[1]),str(x0[2]),str(x0[3]),str(x0[4]),str(x0[5]),\
#                  str(search_range[0]),str(search_range[1]),str(search_range[2]),str(search_range[3]),str(search_range[4]),str(search_range[5]),\
#                  str(T0),str(Tf),str(maxiter),str(maxeval),str(maxaccept),str(dwell),
#                  str(feps),str(retval),schedule,str(boltzmann),str(feval),method,str(quench),str(m),str(n))
#        
#        output = '\t'.join(output)
#        log.write(output)
#        log.write('\n')
#        log.close()

if __name__ == "__main__":
    T = Test()
#    for i in range(11):
#        T.run_test(n_pose=45+i)
    T.run_test(n_pose=50)
    #pylab.show()


